#include <cuda_runtime.h>
#include <stdio.h>
#include <stdlib.h>
#include <time.h>

/* CUDA 错误检查宏 */
#define CHECK_CUDA_ERROR(call) do { \
    cudaError_t err = (call); \
    if (err != cudaSuccess) { \
        fprintf(stderr, "CUDA 错误在 %s:%d - %s\n", __FILE__, __LINE__, cudaGetErrorString(err)); \
        exit(1); \
    } \
} while (0)

/* 复数乘法设备函数 */
__device__ void complex_mul(float2 a, float2 b, float2 *result) {
    result->x = a.x * b.x - a.y * b.y; /* 实部：ac - bd */
    result->y = a.x * b.y + a.y * b.x; /* 虚部：ad + bc */
}

/* CUDA 内核函数：计算 C = A * A^T */
__global__ void matrix_multiply_kernel(float2 *d_A, float2 *d_C, int num_matrices, int rows, int cols) {
    int m = blockIdx.z; /* 当前处理的矩阵索引 */
    if (m >= num_matrices) return;

    int bx = blockIdx.x; /* 子块行索引 */
    int by = blockIdx.y; /* 子块列索引 */
    int tx = threadIdx.x; /* 线程在块内的 x 索引 */
    int ty = threadIdx.y; /* 线程在块内的 y 索引 */

    int i0 = bx * 16; /* 子块起始行 */
    int j0 = by * 16; /* 子块起始列 */

    /* 声明共享内存，加速 A 和 A^T 的访存 */
    __shared__ float2 A_shared[16][32]; /* 存储 A 的子块 */
    __shared__ float2 B_shared[16][32]; /* 存储 A^T 的子块（实际上是 A 的不同行） */

    float2 sum = {0.0f, 0.0f}; /* 每个线程累加的结果 */
    int num_k_blocks = (cols + 31) / 32; /* k 维度分块数，1604/32=50.125，向上取整 51 */

    /* 沿 k 维度分块循环 */
    for (int kb = 0; kb < num_k_blocks; kb++) {
        /* 加载 A_shared：A 的 i0+ty 行，kb*32 到 kb*32+31 列 */
        int i = i0 + ty;
        int k_base = kb * 32;
        if (i < rows) {
            for (int l = 0; l < 2; l++) { /* 每个线程加载 2 个元素 */
                int k = k_base + tx + l * 16;
                if (k < cols) {
                    A_shared[ty][tx + l * 16] = d_A[m * rows * cols + i * cols + k];
                } else {
                    A_shared[ty][tx + l * 16] = {0.0f, 0.0f};
                }
            }
        } else {
            for (int l = 0; l < 2; l++) {
                A_shared[ty][tx + l * 16] = {0.0f, 0.0f};
            }
        }

        /* 加载 B_shared：A 的 j0+ty 行，kb*32 到 kb*32+31 列 */
        int j = j0 + ty;
        if (j < rows) {
            for (int l = 0; l < 2; l++) {
                int k = k_base + tx + l * 16;
                if (k < cols) {
                    B_shared[ty][tx + l * 16] = d_A[m * rows * cols + j * cols + k];
                } else {
                    B_shared[ty][tx + l * 16] = {0.0f, 0.0f};
                }
            }
        } else {
            for (int l = 0; l < 2; l++) {
                B_shared[ty][tx + l * 16] = {0.0f, 0.0f};
            }
        }

        __syncthreads(); /* 同步，确保共享内存加载完成 */

        /* 计算当前子块的部分和 */
        if (i0 + tx < rows && j0 + ty < rows) {
            for (int k = 0; k < 32; k++) {
                float2 prod;
                complex_mul(A_shared[tx][k], B_shared[ty][k], &prod);
                sum.x += prod.x;
                sum.y += prod.y;
            }
        }

        __syncthreads(); /* 同步，准备下一次加载 */
    }

    /* 将结果写入全局内存 */
    if (i0 + tx < rows && j0 + ty < rows) {
        d_C[m * rows * rows + (i0 + tx) * rows + (j0 + ty)] = sum;
    }
}

/* 生成随机复数矩阵 */
void generate_random_matrix(float2 *h_A, int num_matrices, int rows, int cols) {
    int m, i, j;
    for (m = 0; m < num_matrices; m++) {
        for (i = 0; i < rows; i++) {
            for (j = 0; j < cols; j++) {
                int idx = m * rows * cols + i * cols + j;
                h_A[idx].x = (float)rand() / RAND_MAX; /* 实部 */
                h_A[idx].y = (float)rand() / RAND_MAX; /* 虚部 */
            }
        }
    }
}

/* 主函数 */
int main() {
    const int num_matrices = 512 * 3; /* 矩阵总数：1536 */
    const int rows = 50; /* 矩阵行数 */
    const int cols = 1604; /* 矩阵列数 */
    const int runs = 1000; /* 运行次数 */

    size_t size_A = num_matrices * rows * cols * sizeof(float2); /* 输入矩阵大小 */
    size_t size_C = num_matrices * rows * rows * sizeof(float2); /* 输出矩阵大小 */

    /* 分配主机内存 */
    float2 *h_A = (float2 *)malloc(size_A);
    float2 *h_C = (float2 *)malloc(size_C);
    if (h_A == NULL || h_C == NULL) {
        fprintf(stderr, "主机内存分配失败\n");
        exit(1);
    }

    /* 生成随机初始数据 */
    srand((unsigned int)time(NULL));
    generate_random_matrix(h_A, num_matrices, rows, cols);

    /* 分配设备内存 */
    float2 *d_A, *d_C;
    CHECK_CUDA_ERROR(cudaMalloc(&d_A, size_A));
    CHECK_CUDA_ERROR(cudaMalloc(&d_C, size_C));

    /* 创建 CUDA 事件用于计时 */
    cudaEvent_t start_total, stop_total;
    cudaEvent_t start_copy_to, stop_copy_to;
    cudaEvent_t start_kernel, stop_kernel;
    cudaEvent_t start_copy_back, stop_copy_back;
    CHECK_CUDA_ERROR(cudaEventCreate(&start_total));
    CHECK_CUDA_ERROR(cudaEventCreate(&stop_total));
    CHECK_CUDA_ERROR(cudaEventCreate(&start_copy_to));
    CHECK_CUDA_ERROR(cudaEventCreate(&stop_copy_to));
    CHECK_CUDA_ERROR(cudaEventCreate(&start_kernel));
    CHECK_CUDA_ERROR(cudaEventCreate(&stop_kernel));
    CHECK_CUDA_ERROR(cudaEventCreate(&start_copy_back));
    CHECK_CUDA_ERROR(cudaEventCreate(&stop_copy_back));

    /* 定义线程块和网格 */
    dim3 block(16, 16); /* 线程块大小：16x16=256 线程 */
    dim3 grid(4, 4, num_matrices); /* 网格大小：4x4x1536，覆盖所有矩阵和子块 */

    /* 统计变量 */
    float copy_time_sum = 0.0f, kernel_time_sum = 0.0f, total_time_sum = 0.0f;
    float copy_time_max = -1.0f, kernel_time_max = -1.0f, total_time_max = -1.0f;
    float copy_time_min = 1e10f, kernel_time_min = 1e10f, total_time_min = 1e10f;

    /* 运行 1000 次 */
    int run;
    for (run = 0; run < runs; run++) {
        float copy_to_time, kernel_time, copy_back_time, total_time;

        CHECK_CUDA_ERROR(cudaEventRecord(start_total));

        /* 主机到设备拷贝 */
        CHECK_CUDA_ERROR(cudaEventRecord(start_copy_to));
        CHECK_CUDA_ERROR(cudaMemcpy(d_A, h_A, size_A, cudaMemcpyHostToDevice));
        CHECK_CUDA_ERROR(cudaEventRecord(stop_copy_to));

        /* 执行内核 */
        CHECK_CUDA_ERROR(cudaEventRecord(start_kernel));
        matrix_multiply_kernel<<<grid, block>>>(d_A, d_C, num_matrices, rows, cols);
        CHECK_CUDA_ERROR(cudaGetLastError()); /* 检查内核启动错误 */
        CHECK_CUDA_ERROR(cudaEventRecord(stop_kernel));

        /* 设备到主机拷贝 */
        CHECK_CUDA_ERROR(cudaEventRecord(start_copy_back));
        CHECK_CUDA_ERROR(cudaMemcpy(h_C, d_C, size_C, cudaMemcpyDeviceToHost));
        CHECK_CUDA_ERROR(cudaEventRecord(stop_copy_back));

        CHECK_CUDA_ERROR(cudaEventRecord(stop_total));
        CHECK_CUDA_ERROR(cudaEventSynchronize(stop_total));

        /* 计算时间 */
        CHECK_CUDA_ERROR(cudaEventElapsedTime(&copy_to_time, start_copy_to, stop_copy_to));
        CHECK_CUDA_ERROR(cudaEventElapsedTime(&kernel_time, start_kernel, stop_kernel));
        CHECK_CUDA_ERROR(cudaEventElapsedTime(&copy_back_time, start_copy_back, stop_copy_back));
        CHECK_CUDA_ERROR(cudaEventElapsedTime(&total_time, start_total, stop_total));

        float total_copy_time = copy_to_time + copy_back_time;

        /* 更新统计值 */
        copy_time_sum += total_copy_time;
        kernel_time_sum += kernel_time;
        total_time_sum += total_time;

        copy_time_max = total_copy_time > copy_time_max ? total_copy_time : copy_time_max;
        kernel_time_max = kernel_time > kernel_time_max ? kernel_time : kernel_time_max;
        total_time_max = total_time > total_time_max ? total_time : total_time_max;

        copy_time_min = total_copy_time < copy_time_min ? total_copy_time : copy_time_min;
        kernel_time_min = kernel_time < kernel_time_min ? kernel_time : kernel_time_min;
        total_time_min = total_time < total_time_min ? total_time : total_time_min;

        /* 实时打印 */
        printf("运行 %d:\n", run + 1);
        printf("  拷贝时间: %f ms (到设备: %f ms, 从设备: %f ms)\n", total_copy_time, copy_to_time, copy_back_time);
        printf("  CUDA 乘法时间: %f ms\n", kernel_time);
        printf("  总时间: %f ms\n\n", total_time);
    }

    /* 打印统计结果 */
    printf("统计结果（%d 次运行）:\n", runs);
    printf("拷贝时间:\n");
    printf("  平均值: %f ms\n", copy_time_sum / runs);
    printf("  最大值: %f ms\n", copy_time_max);
    printf("  最小值: %f ms\n", copy_time_min);
    printf("CUDA 乘法时间:\n");
    printf("  平均值: %f ms\n", kernel_time_sum / runs);
    printf("  最大值: %f ms\n", kernel_time_max);
    printf("  最小值: %f ms\n", kernel_time_min);
    printf("总时间:\n");
    printf("  平均值: %f ms\n", total_time_sum / runs);
    printf("  最大值: %f ms\n", total_time_max);
    printf("  最小值: %f ms\n");

    /* 释放资源 */
    free(h_A);
    free(h_C);
    CHECK_CUDA_ERROR(cudaFree(d_A));
    CHECK_CUDA_ERROR(cudaFree(d_C));
    CHECK_CUDA_ERROR(cudaEventDestroy(start_total));
    CHECK_CUDA_ERROR(cudaEventDestroy(stop_total));
    CHECK_CUDA_ERROR(cudaEventDestroy(start_copy_to));
    CHECK_CUDA_ERROR(cudaEventDestroy(stop_copy_to));
    CHECK_CUDA_ERROR(cudaEventDestroy(start_kernel));
    CHECK_CUDA_ERROR(cudaEventDestroy(stop_kernel));
    CHECK_CUDA_ERROR(cudaEventDestroy(start_copy_back));
    CHECK_CUDA_ERROR(cudaEventDestroy(stop_copy_back));

    return 0;
}